import os
import logging
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.prompts import PromptTemplate
from llm import get_completion_glm
from pymilvus import MilvusClient
from langchain_core.documents import Document

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class CustomRetriever:
    """自定义检索器，使用MilvusClient直接检索文档"""

    def __init__(self, client, collection_name, embeddings, k=10, output_fields=None):
        self.client = client
        self.collection_name = collection_name
        self.embeddings = embeddings
        self.k = k
        self.output_fields = output_fields or ["text", "metadata"]
        self.title_boost = 1.5  # 标题匹配的权重提升

    def get_relevant_documents(self, query):
        logger.info("正在将查询转换为向量...")
        query_vector = self.embeddings.embed_query(query)
        logger.info(f"查询向量前5维: {query_vector[:5]}")

        logger.info("正在执行向量相似性搜索...")
        try:
            # 首先尝试文件名匹配
            filename_matches = []
            try:
                # 解析查询中的文件名和类型
                query_parts = query.lower().split()
                query_filename = None
                query_doctype = None
                
                for part in query_parts:
                    if '.' in part:  # 可能是文件名
                        query_filename = part
                        query_doctype = part.split('.')[-1].upper()
                
                if query_filename:
                    logger.info(f"检测到文件名查询: {query_filename}, 类型: {query_doctype}")
                    
                    # 获取所有文档的元数据
                    all_docs = self.client.query(
                        collection_name=self.collection_name,
                        filter="",
                        output_fields=["metadata"],
                        limit=15  # 设置检索到的文档数量
                    )
                    
                    # 精确匹配文件名和类型
                    for doc in all_docs:
                        metadata = doc.get("metadata", {})
                        filename = metadata.get("filename", "").lower()
                        doctype = metadata.get("document_type", "")
                        
                        # 精确匹配文件名和类型
                        if filename == query_filename and doctype == query_doctype:
                            matched_doc = self.client.get(
                                collection_name=self.collection_name,
                                ids=[doc["id"]],
                                output_fields=self.output_fields
                            )
                            if matched_doc:
                                filename_matches.extend(matched_doc)
                                logger.info(f"找到精确匹配: {filename} ({doctype})")
                
            except Exception as e:
                logger.warning(f"文件名匹配过程中出现错误: {e}")

            # 如果没有找到精确匹配，执行向量搜索
            if not filename_matches:
                logger.info("未找到精确匹配，执行向量搜索...")
                vector_results = self.client.search(
                    collection_name=self.collection_name,
                    data=[query_vector],
                    limit=self.k,
                    output_fields=self.output_fields,
                )
                final_results = vector_results[0] if vector_results else []
            else:
                final_results = filename_matches

            logger.info(f"Milvus 返回的结果数量: {len(final_results)}")

            documents = []
            for idx, hit in enumerate(final_results):
                logger.info(f"命中 {idx+1}: {hit}")
                entity = hit.get("entity", hit)
                text = entity.get("text", "")
                metadata = entity.get("metadata", {})
                logger.info(f"提取文本长度: {len(text)}, 元数据: {metadata}")

                if text.strip():
                    doc = Document(page_content=text, metadata=metadata)
                    documents.append(doc)
                else:
                    logger.warning("命中项中 text 字段为空，已跳过。")

            logger.info(f"成功转换为 Document 对象数量: {len(documents)}")
            return documents

        except Exception as e:
            logger.error(f"Milvus 查询失败: {e}")
            return []


class ChatbotManager:
    def __init__(
        self,
        model_name: str = "moka-ai/m3e-base",
        device: str = "cpu",
        encode_kwargs: dict = {"normalize_embeddings": True},
        llm_model: str = "glm-4",
        llm_temperature: float = 0.7,
        collection_name: str = "vector_db",
        api_key: str = "EMPTY",
        max_tokens: int = 2048,
    ):
        self.model_name = model_name
        self.device = device
        self.encode_kwargs = encode_kwargs
        self.llm_model = llm_model
        self.llm_temperature = llm_temperature
        self.milvus_db_file = "milvus_data/milvus.db"
        self.collection_name = collection_name
        self.api_key = api_key
        self.max_tokens = max_tokens

        os.makedirs(os.path.dirname(self.milvus_db_file), exist_ok=True)

        logger.info(f"初始化 Embedding 模型: {model_name}，设备: {device}")
        self.embeddings = HuggingFaceEmbeddings(
            model_name=self.model_name,
            model_kwargs={"device": self.device},
            encode_kwargs=self.encode_kwargs,
        )

        self.prompt_template = """
            You are an AI assistant developed by Wuhan University. Your task is to provide accurate and helpful answers based on the provided knowledge base search results.

            Knowledge Base Searching Results:
            {context}

            Additional Metadata:
            - Document Titles: {titles}
            - Document Types: {doc_types}

            User Question: {question}

            Guidelines:
            1. If the question is about a specific document or title, prioritize information from that document.
            2. If the context is empty or insufficient, clearly state what information is missing.
            3. When answering, consider the document structure and hierarchy.
            4. Provide detailed and well-explained answers, citing relevant document sections when possible.
            5. If you're unsure about any information, acknowledge the uncertainty.
            6. Please directly give the answer, without explaining the reasoning process.
            7. All of the above information is the result of searching Wuhan University's knowledge base.
            
            Helpful answer:
        """

        self.client = self._initialize_milvus()

        self.prompt = PromptTemplate(
            template=self.prompt_template,
            input_variables=['context', 'question', 'titles', 'doc_types']
        )

        if self.client:
            self.retriever = CustomRetriever(
                client=self.client,
                collection_name=self.collection_name,
                embeddings=self.embeddings,
                k=5, # 设置返回的文档数量
                output_fields=["text", "metadata"]
            )
        else:
            self.retriever = None
            logger.error("Milvus 初始化失败，无法创建 Retriever")

    def _initialize_milvus(self):
        try:
            logger.info(f"连接 Milvus Lite 本地数据库：{self.milvus_db_file}")
            client = MilvusClient(self.milvus_db_file)

            collections = client.list_collections()
            logger.info(f"已有集合: {collections}")
            if self.collection_name not in collections:
                logger.warning(f"集合 '{self.collection_name}' 不存在，正在创建...")
                dim = 768
                client.create_collection(
                    collection_name=self.collection_name,
                    dimension=dim
                )
                logger.info(f"集合 '{self.collection_name}' 创建成功。")

            logger.info("Milvus 连接成功。")
            return client

        except Exception as e:
            logger.error(f"Milvus 初始化失败: {e}")
            return None
            
    def clear_database(self):
        """
        清空当前向量数据库（删除并重新创建集合）
        
        Returns:
            str: 操作结果消息
        """
        if not self.client:
            return "⚠️ 无法清空数据库，Milvus 客户端未初始化。"
            
        try:
            logger.info(f"正在清空向量数据库，集合名称: {self.collection_name}")
            
            # 检查集合是否存在
            collections = self.client.list_collections()
            if self.collection_name in collections:
                # 删除集合
                self.client.drop_collection(self.collection_name)
                logger.info(f"已删除集合: {self.collection_name}")
                
                # 重新创建集合
                dim = 768
                self.client.create_collection(
                    collection_name=self.collection_name,
                    dimension=dim
                )
                logger.info(f"已重新创建集合: {self.collection_name}")
                
                return "✅ 向量数据库已成功清空！"
            else:
                logger.warning(f"集合 {self.collection_name} 不存在，无需清空。")
                return "⚠️ 向量数据库不存在，无需清空。"
                
        except Exception as e:
            error_msg = f"清空向量数据库时发生错误: {e}"
            logger.error(error_msg)
            return f"⚠️ {error_msg}"

    def get_response(self, query: str) -> str:
        if not self.retriever:
            return "⚠️ 向量数据库未初始化，请检查 Milvus 设置。"

        logger.info(f"接收到用户查询: {query}")

        docs = self.retriever.get_relevant_documents(query)
        logger.info(f"返回的文档数量: {len(docs)}")

        # 提取文档内容、标题、来源和类型
        context = "\n".join([doc.page_content for doc in docs])
        titles = ", ".join(set(doc.metadata.get('title', '') for doc in docs if doc.metadata.get('title')))
        doc_types = ", ".join(set(doc.metadata.get('document_type', '') for doc in docs if doc.metadata.get('document_type')))

        logger.info(f"构建的上下文内容长度: {len(context)}")
        logger.debug(f"上下文内容:\n{context}")

        if not context.strip():
            logger.warning("⚠️ 上下文内容为空，可能未能成功检索到相关内容。")

        formatted_prompt = self.prompt.format(
            context=context,
            question=query,
            titles=titles,
            doc_types=doc_types
        )

        try:
            logger.info("正在向大语言模型发送请求...")
            response = get_completion_glm(
                prompt=formatted_prompt,
                model=self.llm_model,
                temperature=self.llm_temperature,
                api_key=self.api_key,
                max_tokens=self.max_tokens
            )
            logger.info("大语言模型返回响应成功。")
            return response
        except Exception as e:
            logger.error(f"LLM 请求失败: {e}")
            return "⚠️ 无法生成回答，请稍后再试。"
